#!/usr/bin/env python3
"""
Pillow 版推理：python infer.py examples/test_baihe.jpg
"""
import sys
from pathlib import Path
import numpy as np
import tensorflow as tf
from PIL import Image

IMG_SIZE   = 96               # 与训练保持一致
MODEL_PATH = Path('models/herbs_cpu_no_quant.tflite')
CLASS_NAMES = ['baihe', 'dangshen', 'gouqi', 'huaihua', 'jinyinhua']

def load_model():
    if not MODEL_PATH.exists():
        sys.exit(f'❌ 找不到模型：{MODEL_PATH.resolve()}')
    interpreter = tf.lite.Interpreter(model_path=str(MODEL_PATH))
    interpreter.allocate_tensors()
    return interpreter

def preprocess(img_path: Path) -> np.ndarray:
    img = Image.open(img_path).convert('RGB')
    img = img.resize((IMG_SIZE, IMG_SIZE), Image.LANCZOS)
    img = np.array(img, dtype=np.float32)
    # 直接应用 MobileNetV3 的预处理公式
    img = (img / 127.5) - 1.0
    return np.expand_dims(img, 0)


def predict(interpreter, img_array):
    print("输入数据统计:")
    print(f"  范围: [{img_array.min():.3f}, {img_array.max():.3f}]")
    print(f"  shape: {img_array.shape}")
    print(f"  dtype: {img_array.dtype}")

    input_idx = interpreter.get_input_details()[0]['index']
    output_idx = interpreter.get_output_details()[0]['index']

    interpreter.set_tensor(input_idx, img_array)
    interpreter.invoke()
    logits = interpreter.get_tensor(output_idx)[0]

    print(f"原始logits范围: [{logits.min():.3f}, {logits.max():.3f}]")

    prob = tf.nn.softmax(logits).numpy()
    print(f"Softmax后概率分布: {prob}")
    print(f"概率总和: {np.sum(prob):.6f}")

    # 打印所有类别的概率
    print("各类别概率:")
    for i, (name, p) in enumerate(zip(CLASS_NAMES, prob)):
        print(f"  {name}: {p:.2%}")

    idx = int(np.argmax(prob))
    return CLASS_NAMES[idx], float(prob[idx])

if __name__ == '__main__':
    if len(sys.argv) < 2:
        sys.exit('用法：python infer.py <图片路径>')
    img_path = Path(sys.argv[1])
    interpreter = load_model()
    img_array   = preprocess(img_path)
    label, conf = predict(interpreter, img_array)
    print(f'预测：{label}  置信度：{conf:.2%}')